import numpy as np


class StandardScaler:

    def __init__(self):
        self.mean_ = None
        self.scale_ = None

    def fit(self, X):

        """根据训练数据集X获得数据的均值和方差"""
        assert X.ndim == 2, "The dimension of X must be 2"  # 暂时只除了2维数据

        # 批量操作 每一列都计算平均值和标准差
        self.mean_ = np.array([np.mean(X[:, i]) for i in range(X.shape[1])])
        self.scale_ = np.array([np.std(X[:, i]) for i in range(X.shape[1])])

        return self

    def transform(self, X):
        """将X根据这个StandardScaler进行均值方差归一化处理"""
        assert X.ndim == 2, "The dimension of X must be 2"  # 只处理2维 就是矩阵
        assert self.mean_ is not None and self.scale_ is not None, \
            "must fit before transform!"
        assert X.shape[1] == len(self.mean_), \
            "the feature number of X must be equal to mean_ and std_"

        # 防止传进来不是浮点型
        resX = np.empty(shape=X.shape, dtype=float)

        # 遍历每一列
        for col in range(X.shape[1]):
            # 批量操作 这列的每一个数都这么处理
            resX[:, col] = (X[:, col] - self.mean_[col]) / self.scale_[col]
        return resX
